ABC215 G - Colorful Candies 2
問題
提出
やったこと
期待値とかわからないので、分子の値を数え上げ的に求めていく。
（いつにも増して数式の書き方がアレになってしまった。）
色$ c_i までの個数の総和を$ S 、$ c_i までで$ n 個選んだときの期待値の分子を$ X_n と置く。
次に見る色$ c_j の個数を$ k 、$ c_j までで$ n 個選んだときの期待値の分子を$ X'_n とする。
もらうDP的に考えると、$ X'_n は「$ c_i までで$ n 個選び、$ c_j は選ばない」+「$ c_i までで$ n-i 個選び、$ c_j から$ i 個選ぶ」となる。
前者は単に$ X_n 。後者は$ c_i までで$ n-i 個選ぶ各パターンそれぞれに、$ c_j から$ i 個選ぶパターンを組み合わせるので、期待値の分子としては$ \sum_{i=1}^{k} X_{n-i} \times \binom{k}{i} + \binom{S}{n-i}\times\binom{k}{i} 。
（$ n-i < 0 のときは、$ X_{n-i}, \binom{S}{n-i} = 0 とする。）
$ \therefore X'_n = X_n + \sum_{i=1}^{k} X_{n-i}\binom{k}{i} + \binom{S}{n-i}\binom{k}{i}
$ = -\binom{S}{n}+\sum_{i=0}^{k} X_{n-i} \binom{k}{i} + \sum_{i=0}^{k}\binom{S}{n-i}\binom{k}{i}
畳み込みな雰囲気があるので、FPSを使ってみる。
$ f(x) = \sum_{n=0}^{\infin} X_nx^n 、$ f'(x) = \sum_{n=0}^{\infin} X'_nx^n, とすると、↑の式から
$ f'(x) = -(1+x)^S + f(x)(1+x)^k+(1+x)^S(1+x)^k
両辺$ (1+x)^S(1+x)^k で割って
$ f'(x)(1+x)^{-(S+k)} = f(x)(1+x)^{-S} - (1+x)^{-k} + 1
つまり、各色の個数を$ k_i とすると、
$ f(x)(1+x)^{-N} = \sum_i \left( 1-(1+x)^{-k_i} \right)
$ \therefore f(x) = \sum_i \left( (1+x)^N - (1+x)^{N-k_i}\right)
つまり、$ K 個選んだときの期待値は、
$ \frac {\sum_i \left( \binom{N}{K} - \binom{N-k_i}{K} \right)}{\binom{N}{K}}
これをそのまま計算すると$ O(N^2) 。後は公式解説の通り、$ k_i が一致する分をまとめて計算することで$ O(N\sqrt N)になる。
雑感
多項式ゴリ押ししようとすると考察ステップが長くなってしまうので、期待値特有の計算も使えるようになりましょう。